from state import EditorState
from langchain_core.prompts import ChatPromptTemplate
from model.zhipu import get_model
from utils.think_split import get_content

prompt = ChatPromptTemplate.from_template("""
你负责新闻标题的优化
根据用户输入的原始标题{title}生成一个更具吸引力且符合新闻规范的标题(最多二十个字)
除标题外，请不要输出其他内容
""")


def generate_title(state: EditorState):
    title = state['title']
    paragraphs = state['paragraphs']

    model = get_model()
    chains = prompt | model
    response = chains.invoke({"title": title})

    print("============ Title Agent ============")
    print(response.content)

    paragraphs.append("已完成修改的主题:" + response.content)
    return {
        "who": "",
        "title": get_content(response.content),  # 制定标题
        "search": "",  # 搜集信息
        "outlines": {},  # 大纲搜集信息 三级标题
        "paragraphs": paragraphs,  # 段落信息
        "check_idea": {},  # 校对意见
        "comment_idea": {}  # 评论意见
    }

# if __name__ == '__main__':
#     inputs = {
#         "who": "",
#         "title": "中国的近代历史",  # 制定标题
#         "search": "",  # 搜集信息
#         "outlines": {},  # 大纲搜集信息 三级标题
#         "paragraphs": [],  # 段落信息
#         "check_idea": {},  # 校对意见
#         "comment_idea": {}  # 评论意见
#     }
#
#     outputs = generate_title(inputs)
#     print(outputs)
#     print(outputs['title'])
